import os
import warnings
from concurrent.futures import ProcessPoolExecutor, as_completed
from technical import *
from function import *

# 忽略警告信息
warnings.filterwarnings("ignore")

# 设置pandas显示选项
pd.set_option('expand_frame_repr', False)

pd.set_option('display.max_rows', 5000)
c_rate = 1.5 / 10000  # 手续费
t_rate = 1 / 1000  # 印花税
# 后续计算N日后涨跌幅所需参数
day_list = [1, 2, 3, 4, 5, 10, 20]
# 测试时间段，可根据数据时间更改
start_time = '20070101'
end_time = '20241231'
# 文件夹路径
file_path = r'股票历史日线数据/'
# 获取文件夹下的所有csv文件
file_list = [f for f in os.listdir(file_path) if '.csv' in f]

# 读取指数数据
index_df = pd.read_csv('sh000300.csv', encoding='gbk', parse_dates=['candle_end_time'])
index_df.rename(columns={'candle_end_time': '交易日期'}, inplace=True)

def process_file(f):
    # print(f"Processing file: {f}")  # 打印正在处理的文件名
    # 加载数据
    df = load_file(file_path, f)

    # 使用指数数据与股票数据合并，补充停牌的数据
    df = merge_index(df, index_df)
    # 计算交易天数
    df['交易天数'] = df.index + 1
    # 剔除上市交易日期不足N天的股票
    df = df[df['交易天数'] > 250]
    if df.empty:
        return None
    # 计算你需要的技术指标
    df = cal_macd(df)
    # 计算未来表现
    for day in day_list:
        df['%s日后涨跌幅' % day] = (df['收盘价'].shift(0 - day) / df['开盘价'].shift(-1) - 1) * (1 - c_rate - t_rate) * (1 - c_rate)
        df['%s日后_相对指数涨跌幅' % day] = ((df['收盘价'].shift(0 - day) / df['开盘价'].shift(-1) - 1) * (1 - c_rate - t_rate) *
            (1 - c_rate)) - (df['close'].shift(0 - day) / df['open'].shift(-1) - 1)
        df['%s日后是否上涨' % day] = df['%s日后涨跌幅' % day] > 0
        df['%s日后是否上涨' % day].fillna(value=False, inplace=True)
    # 选取指定时间范围内的股票
    df = df[(df['交易日期'] >= pd.to_datetime(start_time)) & (df['交易日期'] <= pd.to_datetime(end_time))]
    df = df[df['成交量'] > 0]
    df = df[df['开盘涨停'].shift(-1) == False]
    # 剔除ST股/退市股
    df['下日_是否ST'] = df['股票名称'].str.contains('ST').shift(-1)
    df['下日_是否S'] = df['股票名称'].str.contains('S').shift(-1)
    df['下日_是否退市'] = df['股票名称'].str.contains('退').shift(-1)
    df = df[df['下日_是否S'] == False]
    df = df[df['下日_是否ST'] == False]
    df = df[df['下日_是否退市'] == False]
    # 删除北交所科创板股票
    df = df[~df['股票代码'].str.contains('bj')]
    df = df[~df['股票代码'].str.contains('sh68')]

    return df

if __name__ == '__main__':
    # 使用ProcessPoolExecutor进行多进程处理
    with ProcessPoolExecutor() as executor:
        futures = {}
        for f in file_list:
            print(f"Submitting file: {f}")  # 打印正在提交的任务文件名
            future = executor.submit(process_file, f)
            futures[future] = f

        dfs = []
        for future in as_completed(futures):
            f = futures[future]
            try:
                df = future.result()
                if df is not None:
                    dfs.append(df)
            except Exception as e:
                print(f"Error processing file {f}: {e}")

    all_df = pd.concat(dfs, ignore_index=True)

    # # 分析数据
    analysis(all_df, day_list)
    all_df = all_df[all_df['signal'] == 1]
